

import * as THREE from 'three';
import { GLTFLoader } from 'three/examples/jsm/Addons.js';

import * as MN from '@tensorflow-models/mobilenet';
import * as tf from '@tensorflow/tfjs-core';

import '@tensorflow/tfjs-backend-cpu';
import '@tensorflow/tfjs-backend-webgl';
import { Utils } from '../src/Utils';


const loader = new GLTFLoader();

const sleep = (delay: number) => new Promise((resolve) => setTimeout(resolve, delay))



export class ModelClassifer {
    scene = new THREE.Scene();
    camera = new THREE.PerspectiveCamera(50, 1, 0.1, 500);
    renderer: THREE.WebGLRenderer;
    model?: MN.MobileNet;
    canvas: HTMLCanvasElement;
    objects = new THREE.Group();
    constructor(canvas: HTMLCanvasElement) {
        this.canvas = canvas;
        const dpr = Math.min(2, window.devicePixelRatio);
        const renderer = new THREE.WebGLRenderer({ canvas, antialias: true });
        renderer.setPixelRatio(dpr);
        this.renderer = renderer;
        let camera = this.camera;
        renderer.setClearColor(0xff, 1);
        camera.position.z = 10;
        camera.position.y = 10;
        camera.lookAt(this.scene.position);
        this.scene.add(new THREE.HemisphereLight(0xaaaaaa, 0x444444, 3));
        const light = new THREE.DirectionalLight(0xffffff, 1.5);
        light.position.set(1, 1, 1);
        this.scene.add(light);
        this.scene.add(this.objects);
        MN.load().then(m => { this.model = m });
    }

    dispose() {
    }

    async add(url: string) {
        const gltf = await loader.loadAsync(url);
        let obj = gltf.scene;
        let box = Utils.computeBoundingBox(obj);
        let size = new THREE.Vector3();
        box.getSize(size);
        let r = Math.max(size.x, size.y, size.z);
        obj.scale.multiplyScalar(3 / r);
        this.objects.clear();
        this.objects.add(obj);
        return await this.render();
    }


    async getTags(logits: tf.Tensor2D) {
        const softmax = tf.softmax(logits);
        const values = await softmax.data();
        softmax.dispose();

        const valuesAndIndices = [];
        for (let i = 0; i < values.length; i++) {
            valuesAndIndices.push({ value: values[i], index: i });
        }
        valuesAndIndices.sort((a, b) => {
            return b.value - a.value;
        });
        const result = [];
        for (let i = 0; i < valuesAndIndices.length; i++) {
            let index = valuesAndIndices[i].index;
            let val = valuesAndIndices[i].value;
            let max = Math.floor(val * 64);
            if (max == 0) {
                break;
            }
            for (let j = max; j > 0; j--) {
                let code = index << 6 | j;
                result.push(code.toString(32));
            }
        }
        return result;
    }


    async render() {
        let tags = new Array<string>();
        if (this.model) {
            for (let i = 0; i < 1; i++) {
                this.renderer.render(this.scene, this.camera);
                this.scene.rotateY(i * Math.PI / 6)
                let t1 = await this.model.infer(this.canvas);
                let logits = <tf.Tensor2D>t1;
                tags = await this.getTags(logits);
                logits.dispose();

                await sleep(50);
            }

        }
        return tags;
    }

    resize() {
        let canvas = this.canvas;
        canvas.width = window.innerWidth;
        canvas.height = window.innerHeight;
        this.camera.aspect = canvas.width / canvas.height;
        this.camera.updateProjectionMatrix();
        this.renderer.setSize(canvas.width, canvas.height);
    }
}
